/*
 * License: Copyright (c) Huawei Technologies Co., Ltd. 2012-2019. All rights reserved.
 * Description: push stream 
 * Date: 2020-01-07 10:20:44
 * LastEditTime : 2020-02-11 11:11:28
 */

#include "dest_engine.h"
#include <iomanip>
#include <cstdio>
#include <sstream>
#include <string>
#include <sys/stat.h>
#include <unistd.h>
#include <sys/time.h>

const int TIMES_SECOND_MICROSECOND = 1000000;
const int BIT_RATE = 1843200;  // 1x1024x1800
const int FRAME_RATE = 30;
const int SLEEP = 10;

HIAI_StatusT DestEngine::Init(const hiai::AIConfig& config, const std::vector<hiai::AIModelDescription>& model_desc)
{
    return HIAI_OK;
}

HIAI_StatusT DestEngine::AvInit(std::shared_ptr<TransData> rawDataInput)
{
    avformat_network_init();
    std::string rtspLink = g_pushStreams[rawDataInput->channelId];
    avformat_alloc_output_context2(&pOutFormatCtx, NULL, "rtsp", rtspLink.c_str());
    if (pOutFormatCtx == NULL) {
        std::cout << "Dest engine allocate output context failed!" << std::endl;
        return HIAI_ERROR;
    }
    pOutStream = avformat_new_stream(pOutFormatCtx, NULL);
    if (pOutStream == NULL) {
        return HIAI_ERROR;
    }
    pOutCodecPar = pOutStream->codecpar;
    pOutCodecPar->codec_type = AVMEDIA_TYPE_VIDEO;
    pOutCodecPar->codec_id = rawDataInput->videoFormat ? AV_CODEC_ID_H265 : AV_CODEC_ID_H264;
    pOutCodecPar->bit_rate = BIT_RATE;
    pOutCodecPar->width = rawDataInput->alignWidth;
    pOutCodecPar->height = rawDataInput->alignHeight;
    pOutCodecPar->codec_tag = 0;
    pOutCodecPar->format = AV_PIX_FMT_YUV420P;

    pOutCodecPar->extradata_size = rawDataInput->len;
    pOutCodecPar->extradata = (uint8_t*)av_malloc(rawDataInput->len + AV_INPUT_BUFFER_PADDING_SIZE);
    if (pOutCodecPar->extradata == NULL) {
        std::cout << "av_malloc failed!"<< std::endl;
        return HIAI_ERROR;
    }
    if (0 != memcpy_s(pOutCodecPar->extradata, rawDataInput->len, rawDataInput->buf.get(), rawDataInput->len)) {
        std::cout << "memcpy failed!"<< std::endl;
        return HIAI_ERROR;
    }

    av_dump_format(pOutFormatCtx, 0, rtspLink.c_str(), 1);
    AVDictionary* format_opts = NULL;
    av_dict_set(&format_opts, "rtsp_transport", "tcp", 0);
    if (avformat_write_header(pOutFormatCtx, &format_opts) != 0) {
        std::cout << "Dest engine write header failed!" << std::endl;
        return HIAI_ERROR;
    }
    isHeaderWrite = true;

    return HIAI_OK;
}

DestEngine::~DestEngine()
{
    if (streamOutputRunner.joinable()) {
        streamOutputRunner.join();
    }
    
    if (pOutFormatCtx && isHeaderWrite) {
        if (av_write_trailer(pOutFormatCtx) != 0) {
            std::cout << "DestEngine failed to write tailer!" << std::endl;
        }
    }
    if (pOutFormatCtx) {
        avio_closep(&pOutFormatCtx->pb);
        avformat_free_context(pOutFormatCtx);
        pOutFormatCtx = NULL;
    }
    std::cout << "DestEngine Closed!" << std::endl;
}

void DestEngine::StreamOutputThread()
{
    while (!stop) {
        dequeMutex.lock();
        if (!transDataQueue.empty()) {
            std::shared_ptr<TransData> streamOutput = transDataQueue.front();
            transDataQueue.pop_front();
            dequeMutex.unlock();

            AVPacket pkt;
            int ret;
            ret = av_new_packet(&pkt, streamOutput->len);
            if (ret != 0) {
                std::cout << "allocate avpacket error!" << std::endl;
                break;
            }
            ret = memcpy_s(pkt.data, streamOutput->len, streamOutput->buf.get(), streamOutput->len);
            if (ret != 0) {
                std::cout << "memcpy error!" << std::endl;
                av_free_packet(&pkt);
                break;
            }
            pkt.size = streamOutput->len;
            pkt.stream_index = pOutStream->index;
            pkt.pts = av_rescale_q(streamOutput->frameId, AVRational{1, FRAME_RATE}, pOutStream->time_base);
            pkt.dts = pkt.pts;
            pkt.duration = av_rescale_q(pkt.duration, pOutStream->time_base, pOutStream->time_base);
            av_interleaved_write_frame(pOutFormatCtx, &pkt);
            av_free_packet(&pkt);
        } else {
            dequeMutex.unlock();
            std::this_thread::sleep_for(chrono::milliseconds(SLEEP));
        }
        
    }
}

HIAI_IMPL_ENGINE_PROCESS("DestEngine", DestEngine, DE_INPUT_SIZE)
{
    std::shared_ptr<TransData> rawDataInput = std::static_pointer_cast<TransData>(arg0);

    rawDataInput->isOutputConnected = isAvInitSuccess ? 1 : 0;

    if (transDataQueue.size() >= 300) {
        cout << "push stream queue size: " << transDataQueue.size() << endl;
    }

    if (SendData(0, "TransData", std::static_pointer_cast<void>(rawDataInput)) != HIAI_OK) {
        std::cout << "SendData of DestEngine error!" << std::endl;
        return HIAI_ERROR;
    }
    if (rawDataInput->eos) {
        stop = true;
        return HIAI_OK;
    }

    if (!isAvInitSuccess) {
        return HIAI_OK;
    }

    // the first data is the SPS and PPS of the video, which is neccessary for a video
    // use this info to create codec
    if (rawDataInput->frameId <= 0) {
        if (HIAI_OK != AvInit(rawDataInput)) {
            isAvInitSuccess = false;
            return HIAI_ERROR;
        }
        isAvInitSuccess = true;
        streamOutputRunner = std::thread(&DestEngine::StreamOutputThread, this);
        return HIAI_OK;
    }
    
    dequeMutex.lock();
    transDataQueue.push_back(rawDataInput);
    dequeMutex.unlock();
    return HIAI_OK;
}
